package org.whale.config;

import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.whale.assistant.DoctorAgent;
import org.whale.dao.ChatMessageRepository;
import org.whale.properties.PGVectorProperties;
import org.whale.properties.SearchProperties;
import org.whale.tool.DoctorTool;

@Configuration
@EnableConfigurationProperties({PGVectorProperties.class,SearchProperties.class})
public class LLMConfig {

    /**
     * 医疗助手智能体
     */
    @Bean
    public DoctorAgent doctorAgent(
            QwenStreamingChatModel chatModel,
            ChatMessageRepository chatMessageStore,
            DoctorTool tool,
            EmbeddingStore<TextSegment> embeddingStore,
            SearchApiWebSearchEngine webSearchEngine,
            EmbeddingModel embeddingModel){

        //内容检索
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(1)
                .minScore(0.8)
                .build();

        return AiServices.builder(DoctorAgent.class)
                //流式对话
                .streamingChatLanguageModel(chatModel)
                //记忆功能
                .chatMemoryProvider((memoryId -> MessageWindowChatMemory.builder().id(memoryId).chatMemoryStore(chatMessageStore).maxMessages(20).build()))
                //调用自定义工具
                .tools(tool,new WebSearchTool(webSearchEngine))
                //RAG知识库检索
                .contentRetriever(contentRetriever)
                .build();
    }

    /**
     * 创建向量存储
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(PGVectorProperties properties, EmbeddingModel embeddingModel) {
        //基于 PgVector的向量存储 - 基于yml配置读取
        return PgVectorEmbeddingStore.builder()
                .table(properties.getTable())
                //.dropTableFirst(true) 每次重启都要重新创建
                .createTable(true)	//自动创建表
                .host(properties.getHost())
                .port(properties.getPort())
                .user(properties.getUser())
                .password(properties.getPassword())
                .dimension(embeddingModel.dimension())	//all-minilm模型的向量维度(简单理解就是内容长度如[111,222 ... 333])
                .database(properties.getDatabase())
                .build();
    }

    /**
     * web搜索引擎
     */
    @Bean
    public SearchApiWebSearchEngine webSearchEngine(SearchProperties properties) {
        return SearchApiWebSearchEngine.builder()
                .engine(properties.getEngine())
                .apiKey(properties.getApiKey())
                .build();
    }
}
